from PyQt5.QtWidgets import QApplication, QMainWindow, QPushButton, QLabel, QVBoxLayout, QWidget, QFileDialog
from PyQt5.QtGui import QPixmap
import sys
import cv2
import numpy as np
from PyQt5.QtCore import Qt
import os
import random
class BirdRecognitionWindow(QMainWindow):
    def __init__(self):
        super().__init__()

        self.setWindowTitle("Bird Recognition")

        self.open_button = QPushButton("Open Image")
        self.open_button.clicked.connect(self.open_image)

        self.open_button2 = QPushButton("tests")
        self.open_button.clicked.connect(self.open_image)

        self.recognize_button = QPushButton("Recognize Bird")
        self.recognize_button.clicked.connect(self.recognize_bird)

        self.batch_button = QPushButton("Batch Identification")
        self.batch_button.clicked.connect(self.batch_identification)

        self.image_label = QLabel()
        self.image_label.setFixedSize(640, 480)

        self.result_label = QLabel("Result: ")

        layout = QVBoxLayout()
        layout.addWidget(self.open_button)
        layout.addWidget(self.open_button2)
        layout.addWidget(self.recognize_button)
        layout.addWidget(self.batch_button)
        layout.addWidget(self.image_label)
        layout.addWidget(self.result_label)

        central_widget = QWidget()
        central_widget.setLayout(layout)
        self.setCentralWidget(central_widget)

    def open_image(self):
        file_name, _ = QFileDialog.getOpenFileName(self, "Open Image", "", "Image Files (*.png *.jpg *.jpeg *.bmp)")
        if file_name:
            pixmap = QPixmap(file_name)
            self.image_label.setPixmap(pixmap.scaled(self.image_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))
            self.image_path = file_name

    def batch_identification(self):
                # 设置输入和输出文件夹
        input_folder = r'./in'
        output_folder = r'./out'

        # 确保输出文件夹存在
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)

        # 列出输入文件夹中的所有文件
        files = os.listdir(input_folder)

        # 初始化您的分类器
        # classifier = Classifier()

        # 遍历所有文件
        for file in files:
            # 检查文件是否是图片
            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp', '.gif')):
                # 读取图片
                img = cv2.imread(os.path.join(input_folder, file))

                # 这里假设有一个函数predict可以预测图片的类别
                # category = classifier.predict(img)

                # 示例：这里我们随机分配一个类别，实际应用中应使用分类器预测
                category = 'blasti' # 假设我们有10个类别

                # 创建类别对应的输出文件夹
                category_folder = os.path.join(output_folder, category)
                if not os.path.exists(category_folder):
                    os.makedirs(category_folder)

                # 将图片保存到相应类别的文件夹
                output_path = os.path.join(category_folder, file)
                cv2.imwrite(output_path, img)
                # self.result_label.setText(f'Saved {file} to {category} folder')

        # print('图片分类完成')
        self.result_label.setText(f"Result: 图片分类完成")

    def recognize_bird(self):
        if hasattr(self, 'image_path') and self.image_path:
            # 这里应该调用你的鸟类识别模型
            # 假设模型返回了一个鸟类名称和置信度
            bird_name, confidence = self._fake_recognize_bird(self.image_path)
            self.result_label.setText(f"Result: {bird_name} (Confidence: {confidence:.2f})")
        else:
            self.result_label.setText("Please load an image first.")

    def _fake_recognize_bird(self, image_path):
        # 这里是模拟的鸟类识别函数
        # 实际应用中，你应该调用一个真实的鸟类识别模型
        return "Mock Bird", 0.95  # 返回一个模拟的鸟类名称和置信度

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = BirdRecognitionWindow()
    window.show()
    sys.exit(app.exec_())